import networkx as nx
import random

def nx_setting(params, nx_type):
    """ 生成网络的点和边

    Args:
        params (_type_): 生成网络的参数
        nx_type (_type_): 生成网络的类型

    Returns:
        _type_: 点和边的数组
    """
    def_dict = {
        "ER": nx_ER,
        "BA": nx_BA,
        "WS": nx_WS,
        "RG": nx_RG,
        "None": nx_None
    }
    G = def_dict[nx_type](*params)
    nodes_array = [] # 数组 存储每个点
    positions = nx.spring_layout(G, scale=500)
    for i in list(G.nodes):
        x, y = positions[i]
        nodes_array.append([x, y])
    edges_array = list(G.edges) # 数组 存储每个边
    return nodes_array, edges_array


def nx_ER(nums, prob):
    G = nx.erdos_renyi_graph(n=nums, p=prob)
    return G

def nx_BA(nums, edges):
    G = nx.barabasi_albert_graph(n=nums, m=edges)
    return G

def nx_WS(nums, neighbors, prob):
    G = nx.watts_strogatz_graph(n=nums, k=neighbors, p=prob)
    return G

def nx_RG(nums, degree):
    G = nx.random_regular_graph(n=nums, d=degree)
    return G

def nx_None(nums):
    G = nx.random_regular_graph(n=nums, d=0)
    return G

if __name__ == "__main__":
    array1, array2 = nx_setting([10, 4, 0.9], 'WS')
    print('ans: ', array1, array2)